from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage


#星火认知大模型Spark Max的URL值，其他版本大模型URL值请前往文档（https://www.xfyun.cn/doc/spark/Web.html）查看
SPARKAI_URL = 'wss://spark-api.xf-yun.com/v3.5/chat'
#星火认知大模型调用秘钥信息，请前往讯飞开放平台控制台（https://console.xfyun.cn/services/bm35）查看
SPARKAI_APP_ID = '8dbb4b5e'
SPARKAI_API_SECRET = 'NzZiZDUwMmFmYzIwZTRiMjVhNzQ0NjNk'
SPARKAI_API_KEY = 'e51a7fb80f7a69d365dcecb6b9a01dc6'
#星火认知大模型Spark Max的domain值，其他版本大模型domain值请前往文档（https://www.xfyun.cn/doc/spark/Web.html）查看
SPARKAI_DOMAIN = 'generalv3.5'

def spark_ask(text):
    # 创建一个ChatSparkLLM对象，用于与SparkAI进行交互
    spark = ChatSparkLLM(
        spark_api_url=SPARKAI_URL,
        spark_app_id=SPARKAI_APP_ID,
        spark_api_key=SPARKAI_API_KEY,
        spark_api_secret=SPARKAI_API_SECRET,
        spark_llm_domain=SPARKAI_DOMAIN,
        streaming=False,
    )
    # 创建一个ChatMessage对象，用于存储用户输入的文本
    messages = [ChatMessage(
        role="user",
        content=text
    )]
    # 创建一个ChunkPrintHandler对象，用于处理SparkAI返回的结果
    handler = ChunkPrintHandler()
    # 调用SparkAI的generate方法，生成结果
    result = spark.generate([messages], callbacks=[handler])
    # print(result)
    # 将结果转换为字符串，并按单引号分割
    input_str = str(result).split('\'')
    # 返回分割后的第一个字符串，即SparkAI返回的结果
    return (input_str[1])

    # input_str = str(result)
    # # 找到 text= 的位置
    # start_index = input_str.find("text='") + len("text='")
    # end_index = input_str.find("'", start_index)
    # # 提取 text= 后面的字符串
    # text_value = input_str[start_index:end_index]
    # print(f'{text_value}')

def main():
    while True:
        question = input("请输入你想问的问题（输入'exit'结束程序）：")
        if question == "exit":
            print("程序已结束。")
            break

        result = spark_ask(question)
        print(">>>"+result)

if __name__ == '__main__':
    main()